package adapters

import (
	"ai-gateway/db"
	"ai-gateway/proto"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

type BaiduStreamAdapter struct{}

func (b *BaiduStreamAdapter) StreamCall(ctx context.Context, req *proto.AIRequest,
	stream proto.AdapterService_StreamCallAIServer) error {
	// 加载百度千帆配置
	var config struct {
		ApiKey     string  `db:"api_key"`
		ApiUrl     string  `db:"api_url"`
		TokenPrice float64 `db:"token_price"`
	}
	err := db.DB.Get(&config, "SELECT api_key, api_url, token_price FROM provider_configs WHERE provider = 'baidu'")
	if err != nil {
		return fmt.Errorf("百度配置加载失败: %v", err)
	}

	// 构造请求（参考百度流式 API 文档）
	requestBody := map[string]interface{}{
		"prompt":     req.Prompt,
		"model":      "ERNIE-Bot", // 替换为实际模型
		"stream":     true,
		"max_tokens": 1024,
	}
	jsonData, _ := json.Marshal(requestBody)

	// 发起流式请求
	httpReq, err := http.NewRequestWithContext(ctx, "POST", config.ApiUrl, strings.NewReader(string(jsonData)))
	if err != nil {
		return fmt.Errorf("请求创建失败: %v", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+config.ApiKey)

	client := &http.Client{}
	resp, err := client.Do(httpReq)
	if err != nil {
		return fmt.Errorf("API 调用失败: %v", err)
	}
	defer resp.Body.Close()

	// 解析 SSE 响应
	reader := resp.Body
	buffer := make([]byte, 1024)
	var fullOutput strings.Builder
	var totalTokens int

	for {
		n, err := reader.Read(buffer)
		if err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("流读取失败: %v", err)
		}

		line := strings.TrimSpace(string(buffer[:n]))
		if line == "data: [DONE]" {
			break
		}
		if !strings.HasPrefix(line, "data: ") {
			continue
		}

		data := strings.TrimPrefix(line, "data: ")
		var baiduResp struct {
			Result struct {
				Text   string `json:"text"`
				Tokens int    `json:"tokens"`
			} `json:"result"`
		}
		if err := json.Unmarshal([]byte(data), &baiduResp); err != nil {
			continue
		}

		// 提取片段
		segment := baiduResp.Result.Text
		totalTokens = baiduResp.Result.Tokens

		if segment != "" {
			fullOutput.WriteString(segment)
			// 推送实时结果到 gRPC 流
			if err := stream.Send(&proto.AIResponse{
				Result:     segment,
				Provider:   "baidu",
				TokenUsed:  int32(totalTokens),
				Cost:       int64(float64(totalTokens) * config.TokenPrice * 100),
				IsFinished: false,
			}); err != nil {
				return fmt.Errorf("流发送失败: %v", err)
			}
		}
	}

	// 发送结束标记
	return stream.Send(&proto.AIResponse{
		Result:     fullOutput.String(),
		Provider:   "baidu",
		TokenUsed:  int32(totalTokens),
		Cost:       int64(float64(totalTokens) * config.TokenPrice * 100),
		IsFinished: true,
	})
}

func init() {
	RegisterStreamAdapter("baidu", &BaiduStreamAdapter{})
}
